#  -*- coding: utf-8 -*-

from strategy.stock_pool.base_stock_pool import BaseStockPool
from factor.factor_module import FactorModule
from data.data_module import DataModule
from util.stock_util import get_trading_dates,get_all_codes_date,get_diff_dates,calc_negative_diff_dates,calc_positive_diff_dates,multi_computer,get_all_codes_trading_date
from pandas import DataFrame,Series
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import pylab as pl
import scipy.signal as signal
import time
from sklearn import linear_model
import pickle
import os
from util.database import DB_CONN
from pymongo import UpdateOne,ASCENDING
"""
趋势策略股票池定义
A:波段1起点（极小值）；
B：波段1终点（极大值）；
C：波段2回调最低点（极小值）；
D：当前日；
"""


class TrendingStockPool(BaseStockPool):

    def multi_get_trending_stock(self,date):
        start_time = time.time()
        all_codes = get_all_codes_trading_date(date)
        args =(date,)
        option_codes_dict = dict()
        option_codes_multi_lst = multi_computer(get_trending_stock_at_date,all_codes,args)
        #print(option_codes_multi_lst)
        for codes_dict in option_codes_multi_lst:
            option_codes_dict.update(codes_dict)



        #特征数据存放入数据库，存入之前先要把同时期的数据删掉
        collection = DB_CONN["trending_strategy_feature"]
        collection.create_index([('code', 1), ('date', 1)])
        collection.create_index([('date', 1)])
        collection.remove({"date":date})

        # 特征数据存放入数据库
        update_requests = []
        for key, value in option_codes_dict.items():
            update_requests.append(
                UpdateOne(
                    {'code': value['code'], 'date': value['date']},
                    {'$set': value},
                    upsert=True)
            )
        # print(update_requests)
        # 批量写入，提高访问效率
        if len(update_requests) > 0:
            start_time = time.time()
            update_result = collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('保存趋势策略特征数据到数据集：%s，日期： %s, 插入：%4d条, 更新：%4d条,耗时：%.3f 秒' %
                  (collection.name, date, update_result.upserted_count, update_result.modified_count,
                   (end_time - start_time)),
                  flush=True)

        self.fiilup_future_result_at_date(date)
        self.fillup_public_feature_data_at_date(date)

        end_time = time.time()
        #print(option_codes_dict)
        print("多进程执行趋势策略选股，日期：%s,候选股票：%s,耗时：%.2f" %(date,option_codes_dict.keys(),(end_time - start_time)))


        return option_codes_dict

    def fiilup_future_result_at_date(self,date):
        # 找出过去的日期的股票列表，计算到目前时间的区间涨幅跌幅,并存入数据库
        dm = DataModule()
        #20天涨跌幅
        past_date = calc_negative_diff_dates(code="999999",is_index=True, date=date, delta_days=-50)
        collection = DB_CONN["trending_strategy_option_stocks"]
        past_date_data = collection.find_one({'date': past_date})
        if past_date_data is None:
            print(f"{past_date}:没有数据！")
            return
        codes = past_date_data['option_stocks']
        update_requests = []
        for code in codes:
            print(f"查找 {past_date}:{code}特征数据！")

            result_dict = dict()
            buy_date = calc_positive_diff_dates(code=code, date=past_date, delta_days=1)

            code_df = dm.get_k_data(code, index=False, autype="hfq", begin_date=buy_date, end_date=date)
            if code_df.index.size == 0:
                continue
            code_df.set_index(['date'], inplace=True)
            #print(code_df)
            #先找整个区间最大值，再找起始日到最大日之间的最小值
            close_ser = code_df["close"]
            #beg_min = code_df.loc[past_date]['low']
            #买入日的开盘价
            beg_open = code_df.loc[buy_date]['open']

            max_value = close_ser.max()
            max_idx = close_ser.idxmax()
            max_diff_dates = get_diff_dates(code,buy_date,max_idx)
            #print("maxidx: %s,max_value: %s" %(max_idx,max_value))
            code_to_max_df = code_df.loc[buy_date:max_idx]
            min_value = code_to_max_df['close'].min()
            min_idx = code_to_max_df['close'].idxmin()
            min_diff_dates =  get_diff_dates(code,buy_date,min_idx)
            #print("min_idx: %s,min_value: %s" %(min_idx,min_value))

            max_change_rate = round(100 * (max_value - beg_open)/beg_open,2)
            min_change_rate = round(100 * (min_value - beg_open)/beg_open,2)

            # #print(diff_dates,max_change_rate,min_value,beg_min)
            #最大值日2个交易日（保证可卖出），最大值收益超过8%，最小值不低于起始日的最低价（大阳底部）
            # if diff_dates >= 1 and max_change_rate >= 8  and min_value >= beg_min :
            #     code_feature_df.loc[i,"Y"] = 1
            result_dict['max_change_rate'] = max_change_rate
            result_dict['min_change_rate'] = min_change_rate
            result_dict['max_diff_dates'] = max_diff_dates
            result_dict['min_diff_dates'] = min_diff_dates

            update_requests.append(
                UpdateOne(
                    {'code': code, 'date': past_date},
                    {'$set': result_dict},
                    upsert=True)
            )
        # 批量写入，提高访问效率
        collection = DB_CONN["trending_strategy_feature"]
        if len(update_requests) > 0:
            start_time = time.time()
            update_result = collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('更新趋势策略特征数据到数据集：%s，日期： %s, 插入：%4d条, 更新：%4d条,耗时：%.3f 秒' %
                  (collection.name, past_date, update_result.upserted_count, update_result.modified_count,
                   (end_time - start_time)),
                  flush=True)
        return

    def fillup_public_feature_data_at_date(self,date):
        #补充大盘公共数据到个股特征里

        public_dict = dict()

        # 当日涨跌停数据(涨跌幅大于9.9或小于-9.9)


        return

    def export_feature_data(self):
        feature_collection = DB_CONN[f'{self.strategy_name}_strategy_feature']
        daily_cursor = feature_collection.find(
            {'date': {'$gte': self.begin_date, '$lte': self.end_date}},
            sort=[('date', ASCENDING)],
            projection={'_id': False})
        feature_data_df = DataFrame([x for x in daily_cursor])
        feature_data_df.set_index("code", inplace=True)
        feature_data_df.to_csv(f"{self.stock_pool_path}/{self.strategy_name}_{self.begin_date}_{self.end_date}_feature.csv",encoding="UTF-8_sig")
        return

    def get_option_stocks(self):
        """
        实现股票池选股逻辑，找到指定日期范围的候选股票
        条件：0 < PE < 30, 按从小到大排序，剔除停牌后，取前100个；再平衡周期：7个交易日
        :return: tuple，再平衡的日期列表，以及一个dict(key: 再平衡日, value: 当期的股票列表)
        """

        dm = DataModule()

        # 获取日期范围的交易日历
        all_dates = get_trading_dates(self.begin_date, self.end_date)

        # 缓存股票和其对应有交易的日期
        code_dates_cache = dict()

        # 调整日和其对应的股票
        rebalance_date_codes_dict = dict()
        rebalance_dates = []

        # 保存上一期的股票池
        last_phase_codes = []
        # 所有的交易日数
        dates_count = len(all_dates)

        # 用再平衡周期作为步长循环
        update_requests = []
        for index in range(0, dates_count, self.interval):
            # 当前的调整日
            rebalance_date = all_dates[index]
            print(rebalance_date)
            self.logger.info('调整日：' + rebalance_date)

            # 获取本期符合条件的备选股票
            this_phase_option_codes_dict = self.multi_get_trending_stock(rebalance_date)
            #this_phase_option_codes = get_trending_stock_at_date(None,rebalance_date)
            this_phase_option_codes = list(this_phase_option_codes_dict.keys())
            option_codes_copy = this_phase_option_codes.copy()
            self.logger.info('原始备选股票')
            this_phase_option_codes.sort()
            self.logger.info(this_phase_option_codes)

            # 本期入选的股票代码列表
            this_phase_codes = []

            # 找到在上一期的股票池，但是当前停牌的股票，保留在当期股票池中
            if len(last_phase_codes) > 0:
                print("上一期股票池个数：%4d" % len(last_phase_codes))
                for code in last_phase_codes:
                    if code not in list(code_dates_cache.keys()):
                        daily_ks = dm.get_k_data(code, autype=None,index = False, begin_date=self.begin_date, end_date=self.end_date)
                        daily_ks.set_index(['date'], inplace=True)
                        # 只保留交易日的数据
                        daily_ks = daily_ks[daily_ks['is_trading'] == True]
                        code_dates_cache[code] = list(daily_ks.index)

                    if rebalance_date not in code_dates_cache[code]:
                        this_phase_codes.append(code)


            self.logger.info('上期停牌的股票')
            this_phase_codes.sort()
            self.logger.info(this_phase_codes)

            # 剩余的位置用当前备选股票的,先去重
            for dup_code in this_phase_codes:
                if dup_code in this_phase_option_codes:
                    this_phase_option_codes.remove(dup_code)

            option_size = len(this_phase_option_codes)
            if option_size > (100 - len(this_phase_codes)):
                this_phase_codes += this_phase_option_codes[0:100 - len(this_phase_codes)]

            else:
                this_phase_codes += this_phase_option_codes

            # 当期股票池作为下次循环的上期股票池
            last_phase_codes = this_phase_codes

            # 保存到返回结果中
            rebalance_date_codes_dict[rebalance_date] = this_phase_codes
            rebalance_dates.append(rebalance_date)

            self.logger.info('当前最终的备选票：%s' % rebalance_date)
            this_phase_codes.sort()
            self.logger.info(this_phase_codes)

            #存入数据库
            update_requests.append(
                UpdateOne(
                    {'date': rebalance_date},
                    {'$set': {"date":rebalance_date,
                              "stocks":rebalance_date_codes_dict[rebalance_date],
                              "option_stocks":option_codes_copy}},
                    upsert=True)
            )

        # 批量写入，提高访问效率
        collection = DB_CONN["trending_strategy_option_stocks"]
        collection.create_index([('date', 1)])
        if len(update_requests) > 0:
            start_time = time.time()
            update_result = collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('保存趋势策略特征数据到数据集：%s，插入：%4d条, 更新：%4d条,耗时：%.3f 秒' %
                  (collection.name, update_result.upserted_count, update_result.modified_count,(end_time-start_time)),
                  flush=True)

        return rebalance_dates, rebalance_date_codes_dict

def insert_feature_data_at_date(code,date):

    feature_dict = dict()
    fm = FactorModule()

    #rs数据
    date_rs_df = fm.get_single_stock_factors(code, "rs", False, date, date)
    if len(date_rs_df) == 1 and "rs_250" in date_rs_df.columns.values:
        feature_dict["date_C_rs_60"] = date_rs_df.loc[0]['rs_60']
        feature_dict["date_C_rs_120"] = date_rs_df.loc[0]['rs_120']
        feature_dict["date_C_rs_250"] = date_rs_df.loc[0]['rs_250']


    return feature_dict


def get_trending_stock_at_date(codes,date):
    start_time = time.time()
    option_codes_dict = dict()
    fm = FactorModule()
    if codes == None:
        codes = ['000860',"300628"]

    dm = DataModule()
    change_rate_minimum = 4

    for code in codes:

        try:
            #每次循环需要建一个新的对象，否则赋值后下一个循环会被覆盖掉，从而影响保存的值
            code_data_s = dict()

            print("开始计算股票：%s的趋势" % code)
            # 取近3年数据
            begin_date = calc_negative_diff_dates(code=None,date = date,delta_days = -750)
            if begin_date == False:
                begin_date = None

            code_df = dm.get_k_data(code,index=False,autype="hfq",begin_date=begin_date,end_date=date)

            if code_df.index.size == 0:
                continue

            code_df.set_index(['date'], inplace=True)
            ma_20,ma_250,change_rate_above_ma20 = None,None,None

            #第一步，找出当日涨幅>4%，收盘价大于MA20且超过MA20在10%以内,收盘价大于MA250的数据,不足1年或不足20日的新股次新股不分析
            change_rate_df = fm.get_single_stock_factors(code,"change_rate",False,date,date)
            hfq_ma_df = fm.get_single_stock_factors(code,"hfq_ma",False,date,date)
            vol_ratio_df = fm.get_single_stock_factors(code,"vol_ratio",False,date,date)
            close = code_df.loc[date]['close']
            amount = code_df.loc[date]['amount']
            if len(change_rate_df) == 1:
                change_rate = change_rate_df.loc[0]['close_change_rate']
            if len(hfq_ma_df) == 1:
                ma_20 = hfq_ma_df.loc[0]['ma20']
                ma_250 = hfq_ma_df.loc[0]['ma250']
                if pd.isnull(ma_20) or pd.isnull(ma_250):
                    print("S1:没有20日线或250日线，为新股或次新股")
                    continue
                change_rate_above_ma20 = round(100*(close-ma_20)/ma_20,2)
            if len(vol_ratio_df) == 1:
                vol_ratio_5 = vol_ratio_df.loc[0]["vol_ratio_5"]
                vol_ratio_10 = vol_ratio_df.loc[0]["vol_ratio_10"]
                if pd.isnull(vol_ratio_5) or pd.isnull(vol_ratio_10):
                    print("S1:没有5日或10日量比")
                    continue

            #涨跌幅四舍五入
            print("S1:当日涨幅：%.2f,收盘价：%.2f,ma20：%.2f,ma250: %.2f,相对ma20涨幅：%.2f,成交额：%.2f 亿元,5日量比：%.2f，10日量比：%.2f"
                  % (change_rate,close,ma_20,ma_250,change_rate_above_ma20,amount/1E8,vol_ratio_5,vol_ratio_10))
            if round(change_rate,1) < change_rate_minimum\
                or close < ma_20\
                or close < ma_250\
                or change_rate_above_ma20 > 10\
                or amount < 1E8\
                or vol_ratio_5 < 1\
                or vol_ratio_10 < 1:
                continue

            #第二步，找出最近的极大值极小值点,计算波段1和波段2涨跌幅，
            #波段1涨幅>50,波段2回调<30
            #最近5天的20日线拟合斜率向上
            #C点到D点的拟合斜率向上

            extrem_date_index = []
            extrem_date_values = []
            close_array = code_df["close"]
            code_df['ma20'] = code_df["close"].rolling(20).mean()
            close_np = close_array.values

            extrem_max_index = signal.argrelextrema(close_np, np.greater, order=30)[0]
            extrem_min_index = signal.argrelextrema(-close_np, np.greater, order=30)[0]

            extrem_max_date_index = code_df.index[extrem_max_index]
            extrem_min_date_index = code_df.index[extrem_min_index]

            #波段2至少回调50天，所以波段2期间50天内的极大值点舍去
            [extrem_date_index.append(x) for x in extrem_max_date_index]
            [extrem_date_index.append(x) for x in extrem_min_date_index]
            #降序排序,值的顺序与排序后的日期顺序相同
            extrem_date_index.sort(reverse = True)
            [extrem_date_values.append(code_df.loc[x]['close']) for x in extrem_date_index]

            #print(extrem_date_index)
            date_A,date_A_value = None,None
            date_B,date_B_value = None,None
            date_C,date_C_value = None,None
            for extrem_date in extrem_date_index:
                #找到第一个大于90日以上间隔的极大值，在此范围内找最大值，最大值日间隔大于40天，那么就是B
                if date_B == None and extrem_date in extrem_max_date_index:
                    diff_dates = get_diff_dates(code,extrem_date,date)
                    if diff_dates > 90:
                        #list切片是左闭右开，所以要+1
                        date_max_value_least_90 = max(extrem_date_values[0:extrem_date_index.index(extrem_date)+1])
                        date_max_least_90 = extrem_date_index[extrem_date_values.index(date_max_value_least_90)]
                        if get_diff_dates(code,date_max_least_90,date) > 40:
                            date_B = date_max_least_90
                            date_B_value = extrem_date_values[extrem_date_index.index(date_B)]

                #找到B之后，B之前的最小值是C（如果有多个相同的最小值，取最早的一个）
                if date_B is not None:
                    date_C_value = code_df.loc[date_B:date,'close'].min()
                    #date_C_value = min(extrem_date_values[0:extrem_date_index.index(date_B)+1])
                    date_C = code_df[(code_df.index >date_B) & (code_df["close"] == date_C_value)].index.tolist()[0]

                    #找到B之后，B之后的第一个极小值就是A
                    if date_A == None:
                        for min_date in extrem_date_index[extrem_date_index.index(date_B):]:
                            if min_date in extrem_min_date_index:
                                date_A = min_date
                                date_A_value = extrem_date_values[extrem_date_index.index(date_A)]
                                break


            #A是在找到B的前提下找到的，所以A不为空，其他都不是空
            if date_A is None:
                print("S2:没找到date_A,不符合趋势形态")
                continue

            change_rate_AB,diff_date_AB,change_rate_BC,diff_date_BC = 0,0,0,0
            change_rate_AB = round(100*(date_B_value - date_A_value)/date_A_value,2)
            diff_date_AB = get_diff_dates(code,date_A,date_B)
            change_rate_BC = round(100 * (date_C_value - date_B_value) / date_B_value, 2)
            diff_date_BC = get_diff_dates(code, date_B, date_C)
            change_rate_BD = round(100*(close - date_B_value) / date_B_value, 2)
            diff_date_BD = get_diff_dates(code,date_B,date)
            print("S2:A日：%s,B日：%s,C日：%s,AB涨幅：%.2f,AB间隔：%s天，BC回调：%.2f，BC间隔：%s天，BD涨幅: %.2f,BD间隔：%.2f"
                   % (date_A,date_B,date_C,change_rate_AB,diff_date_AB,change_rate_BC,diff_date_BC,change_rate_BD,diff_date_BD))

            clf = linear_model.LinearRegression()
            date_near_4 = calc_negative_diff_dates(code,False,date,-4)
            #print(date_near_4)
            clf.fit(np.arange(5).reshape(-1,1),code_df['ma20'].loc[date_near_4:date])
            k_20 = clf.coef_
            print("S2:最近5日的20日线拟合的斜率：%.2f" % clf.coef_)
            clf.fit(np.arange(get_diff_dates(code,date_C,date)).reshape(-1,1),code_df['close'].loc[date_C:date])
            k_CD = clf.coef_
            print("S2:从C点到当日日线收盘价拟合的斜率：%.2f" % clf.coef_)


            if change_rate_AB < 50 or change_rate_BC < -30 or change_rate_BC > 0\
                or k_20 < -0.1 or k_CD < -0.1 \
                or change_rate_BD > 5:
                print("S2:形态不满足量化条件")
                continue

            #第三步，成交量判断
            #波段1成交量大，波段2成交量小
            amount_mean_AB= code_df.loc[date_A:date_B,"amount"].mean()
            amount_mean_BD= code_df.loc[date_B:date,"amount"].mean()
            change_rate_amount_AB_BC = round(100 * (amount_mean_AB - amount_mean_BD)/amount_mean_AB,2)
            print("S3:波段1(AB)平均成交额：%.2f亿元,波段2(BD)平均成交额：%.2f亿元,差值比例: %.2f"
                  %(amount_mean_AB/1E8,amount_mean_BD/1E8,change_rate_amount_AB_BC))
            if amount_mean_AB < amount_mean_BD :
                print("S3: AB量能比BD低，不满足波段1放量，波段2缩量的条件")
                #continue
            code_data_s["date"] = date
            code_data_s["code"] = code
            code_data_s["date_A"] = date_A
            code_data_s["date_B"] = date_B
            code_data_s["date_C"] = date_C

            #插入其他一些辅助的特征数据
            feature_dict = insert_feature_data_at_date(code,date)
            code_data_s.update(feature_dict)

            option_codes_dict[code] = code_data_s

            #作图
            # code_df.plot(title='close', y=['close',"ma20"], kind='line')
            # plt.plot(extrem_max_index, close_np[extrem_max_index], 'o')
            # plt.plot(extrem_min_index, close_np[extrem_min_index], '+')
            # for x, y in zip(extrem_max_index, close_np[extrem_max_index]):
            #     plt.text(x, y + 0.3, '%s' % (code_df.index[x]), ha='center', va='bottom', fontsize=10.5)
            # for x, y in zip(extrem_min_index, close_np[extrem_min_index]):
            #     plt.text(x, y + 0.3, '%s' % (code_df.index[x]), ha='center', va='bottom', fontsize=10.5)
            # plt.show()
        except:
            print("股票 %s,计算趋势指标是错误，日期：%s" % (code,date))

    end_time = time.time()
    print("日期：%s,候选股票：%s,耗时：%.2f"% (date,option_codes_dict.keys(),(end_time - start_time)))
    return option_codes_dict




if __name__ == '__main__':
    tp = TrendingStockPool(strategy_name="trending", begin_date="2015-01-01",end_date="2018-11-26",interval=1)

    #tp.multi_get_trending_stock("2018-10-22")
    #print(tp.get_option_stocks())
    #result = get_trending_stock_at_date(codes = None,date="2018-10-22")
    #print(result[1])


